package com.zzz.thrift.gateway.service;

import com.alibaba.cloud.nacos.NacosDiscoveryProperties;
import com.alibaba.cloud.nacos.NacosServiceManager;
import com.alibaba.fastjson.JSONObject;
import com.alibaba.nacos.api.naming.NamingService;
import com.alibaba.nacos.api.naming.pojo.Instance;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import org.apache.thrift.TApplicationException;
import org.apache.thrift.TException;
import org.apache.thrift.TProcessor;
import org.apache.thrift.protocol.TCompactProtocol;
import org.apache.thrift.protocol.TMessage;
import org.apache.thrift.protocol.TMessageType;
import org.apache.thrift.protocol.TProtocol;
import org.apache.thrift.transport.TSocket;
import org.apache.thrift.transport.TTransport;
import org.apache.thrift.transport.layered.TFastFramedTransport;
import org.springframework.beans.BeansException;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;
import org.springframework.stereotype.Component;

import java.util.ArrayList;
import java.util.List;
import java.util.Random;

@Component
@Slf4j
public class ThriftProxyNacosProcessor implements TProcessor, ApplicationContextAware {

    private ApplicationContext applicationContext;

    @Override
    public void setApplicationContext(ApplicationContext applicationContext) throws BeansException {
        this.applicationContext = applicationContext;
    }

    @Autowired
    private NacosServiceManager nacosServiceManager;

    @Autowired
    private NacosDiscoveryProperties nacosDiscoveryProperties;

    @SneakyThrows
    @Override
    public void process(TProtocol in, TProtocol out) throws TException {
        TMessage msg = in.readMessageBegin();
        String applicationName = msg.name.split(":")[0];
        String serviceName = msg.name.split(":")[1];
        List<ProxyStruct> inDatas = ThriftMessageUtils.readData(in);
        in.readMessageEnd();

        //转发请求到目标thrift服务端
        //应用实例列表
        NamingService namingService = nacosServiceManager.getNamingService(nacosDiscoveryProperties.getNacosProperties());
        List<Instance> instances = namingService.getAllInstances(applicationName);
        List<Instance> avaliableInstances = new ArrayList<>();
        for (Instance instance : instances) {
            if (instance.isHealthy()) {
                avaliableInstances.add(instance);
            }
        }
        if (avaliableInstances.isEmpty()) {
            throw new TApplicationException(TApplicationException.INVALID_PROTOCOL, applicationName + " not found");
        }
        //简单实现负载均衡（随机算法）
        Instance instance = getInstance(avaliableInstances);
        String serviceInfoJson = instance.getMetadata().get("server.info");
        if (serviceInfoJson == null) {
            throw new TApplicationException(TApplicationException.INVALID_PROTOCOL, applicationName + " not found");
        }
        List<String> serviceInfos = JSONObject.parseArray(serviceInfoJson, String.class);
        int port = 0;
        for (String serviceInfo : serviceInfos) {
            String[] service = serviceInfo.split(":");
            if (service[0].equals(serviceName)) {
                port = Integer.valueOf(service[1]);
                break;
            }
        }
        if (port == 0) {
            throw new TApplicationException(TApplicationException.INVALID_PROTOCOL, serviceName + " not found");
        }
        //创建连接
        TSocket socket = new TSocket(instance.getIp(), port);
        socket.setTimeout(3000);
        TTransport transport = socket;
        transport = new TFastFramedTransport(transport);
        TProtocol tProtocol = new TCompactProtocol(transport);
        if (!transport.isOpen()) {
            transport.open();
        }
        //发送消息
        msg = new TMessage(msg.name.substring(applicationName.length() + 1), msg.type, msg.seqid);
        tProtocol.writeMessageBegin(msg);
        ThriftMessageUtils.wirteData(tProtocol, inDatas);
        tProtocol.writeMessageEnd();
        tProtocol.getTransport().flush();
        // 读取返回结果
        int seqid = msg.seqid;
        String methodName = msg.name;
        msg = tProtocol.readMessageBegin();
        if (msg.type == TMessageType.EXCEPTION) {
            TApplicationException x = new TApplicationException();
            x.read(tProtocol);
            tProtocol.readMessageEnd();
            throw x;
        }
        if (msg.seqid != seqid) {
            throw new TApplicationException(TApplicationException.BAD_SEQUENCE_ID, methodName + " failed: out of sequence response");
        }
        inDatas = ThriftMessageUtils.readData(tProtocol);
        tProtocol.readMessageEnd();

        out.writeMessageBegin(msg);
        ThriftMessageUtils.wirteData(out, inDatas);
        out.writeMessageEnd();
        out.getTransport().flush();
    }

    /**
     * 实现负载均衡（随机算法）
     * @param instances
     * @return
     */
    public Instance getInstance(List<Instance> instances) {
        Random random = new Random();
        int randomPos = random.nextInt(instances.size());
        return instances.get(randomPos);
    }

}
